package com.gitee.c0dehvb.reactive.sample.r2dbc.connection;

import com.gitee.c0dehvb.reactive.sample.r2dbc.ShardingRewriteEngine;
import io.r2dbc.spi.Batch;
import io.r2dbc.spi.Connection;
import io.r2dbc.spi.Statement;

/**
 * 分片连接
 * @author LiYilin
 * @since 2022/6/15 9:36 PM
 **/
public class ShardingConnection extends ConnectionDelegator {
    private final ShardingRewriteEngine shardingRewriteEngine;

    public ShardingConnection(Connection connection, ShardingRewriteEngine shardingRewriteEngine) {
        super(connection);
        this.shardingRewriteEngine = shardingRewriteEngine;
    }

    @Override
    public Batch createBatch() {
        return new ShardingBatch(connection.createBatch(), shardingRewriteEngine);
    }

    @Override
    public Statement createStatement(String sql) {
        return new ShardingStatement(sql, connection, shardingRewriteEngine);
    }
}
